import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain.agents import tool
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.utils.function_calling import convert_to_openai_function
from langchain.agents.format_scratchpad import format_to_openai_function_messages
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain.schema.agent import AgentFinish
import pymysql
from loguru import logger
# 1、加载配置环境
load_dotenv()

# 2.初始化模型
qwen = ChatOpenAI(
    model="qwen-max",
    api_key="sk-079868c30081417fb54343c6919ba60d",
    openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0
)

# 3、定义提示模板
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "你是一个专业的数据分析师，请先查询知识库获取专业知识和数据库表字段元数据，再根据用户意图决定是否查询数据库，切记sql必须指明库名database.table ，最后回答用户问题"
        ),
        ("user", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ]
)

# 4.定义工具
@tool
def execute_sql_tool(sql):
    """查询数据库，执行一条sql语句"""
    db_config = {
        'user': os.getenv('DB_USER', 'root'),
        'password': os.getenv('DB_PASSWORD', '000000'),
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 3306)),
        'database': os.getenv('DB_NAME', 'test'),
        'charset': os.getenv('DB_CHARSET', 'utf8')
    }
    try:
        connection = pymysql.connect(**db_config)
        logger.info(f"执行sql语句：{sql}")
        with connection.cursor() as cursor:
            cursor.execute(sql)
            # 根据SQL类型处理结果
            if sql.strip().lower().startswith('select'):
                result = cursor.fetchall()
                columns = [desc[0] for desc in cursor.description]
                return [dict(zip(columns, row)) for row in result]
            else:
                connection.commit()
                return {"message": "操作成功", "rows_affected": cursor.rowcount}
    except pymysql.Error as e:
        return {"error": f"数据库错误: {str(e)}", "sql": sql}
    except Exception as e:
        return {"error": f"未知错误: {str(e)}", "sql": sql}
    finally:
        if connection:
            connection.close()
@tool
def bank_knowledge_tool() -> dict:
    """查询银行数据库知识库，包含表结构和字段元数据"""
    knowledge = {
        "银行数据库": {
            "库名称": "bank_db",
            "tables": {
                "客户信息表": {
                    "表名称": "customers",
                    "字段": {
                        "customer_id": {"含义": "客户唯一标识", "计算口径": "自增整数"},
                        "customer_name": {"含义": "客户姓名", "计算口径": " varchar(100)"},
                        "id_number": {"含义": "身份证号码", "计算口径": " varchar(18)"},
                        "phone": {"含义": "联系电话", "计算口径": " varchar(20)"},
                        "email": {"含义": "电子邮箱", "计算口径": " varchar(50)"},
                        "address": {"含义": "联系地址", "计算口径": " varchar(200)"},
                        "create_date": {"含义": "客户创建日期", "计算口径": " datetime"}
                    }
                },
                "账户表": {
                    "表名称": "accounts",
                    "字段": {
                        "account_id": {"含义": "账户唯一标识", "计算口径": "自增整数"},
                        "customer_id": {"含义": "关联客户ID", "计算口径": "外键关联customers表customer_id"},
                        "account_number": {"含义": "账号", "计算口径": " varchar(30)"},
                        "account_type": {"含义": "账户类型", "计算口径": " enum('储蓄账户', '支票账户', '信用卡账户')"},
                        "balance": {"含义": "账户余额", "计算口径": " decimal(15,2)"},
                        "open_date": {"含义": "开户日期", "计算口径": " datetime"},
                        "status": {"含义": "账户状态", "计算口径": " enum('正常', '冻结', '注销')"}
                    }
                },
                "交易记录表": {
                    "表名称": "transactions",
                    "字段": {
                        "transaction_id": {"含义": "交易唯一标识", "计算口径": "自增整数"},
                        "account_id": {"含义": "交易账户ID", "计算口径": "外键关联accounts表account_id"},
                        "transaction_type": {"含义": "交易类型", "计算口径": " enum('存款', '取款', '转账', '支付')"},
                        "amount": {"含义": "交易金额", "计算口径": " decimal(15,2)"},
                        "transaction_date": {"含义": "交易日期时间", "计算口径": " datetime"},
                        "counterparty": {"含义": "交易对手", "计算口径": " varchar(100)"},
                        "description": {"含义": "交易描述", "计算口径": " varchar(200)"},
                        "transaction_status": {"含义": "交易状态", "计算口径": " enum('成功', '失败', '处理中')"}
                    }
                },
                "贷款表": {
                    "表名称": "loans",
                    "字段": {
                        "loan_id": {"含义": "贷款唯一标识", "计算口径": "自增整数"},
                        "customer_id": {"含义": "客户ID", "计算口径": "外键关联customers表customer_id"},
                        "loan_amount": {"含义": "贷款金额", "计算口径": " decimal(15,2)"},
                        "loan_type": {"含义": "贷款类型", "计算口径": " enum('个人贷款', '房贷', '车贷', '企业贷款')"},
                        "interest_rate": {"含义": "贷款利率", "计算口径": " decimal(5,2)"},
                        "start_date": {"含义": "贷款开始日期", "计算口径": " datetime"},
                        "end_date": {"含义": "贷款结束日期", "计算口径": " datetime"},
                        "status": {"含义": "贷款状态", "计算口径": " enum('还款中', '已还清', '逾期', '违约')"},
                        "remaining_balance": {"含义": "剩余还款金额", "计算口径": " decimal(15,2)"}
                    }
                },
                "理财产品表": {
                    "表名称": "financial_products",
                    "字段": {
                        "product_id": {"含义": "产品唯一标识", "计算口径": "自增整数"},
                        "product_name": {"含义": "产品名称", "计算口径": " varchar(100)"},
                        "product_type": {"含义": "产品类型", "计算口径": " enum('基金', '债券', '保险', '信托')"},
                        "risk_level": {"含义": "风险等级", "计算口径": " enum('R1', 'R2', 'R3', 'R4', 'R5')"},
                        "expected_return": {"含义": "预期收益率", "计算口径": " decimal(5,2)"},
                        "investment_period": {"含义": "投资期限", "计算口径": " varchar(50)"},
                        "min_investment": {"含义": "最低投资额", "计算口径": " decimal(15,2)"},
                        "start_date": {"含义": "开始日期", "计算口径": " datetime"},
                        "end_date": {"含义": "结束日期", "计算口径": " datetime"},
                        "status": {"含义": "产品状态", "计算口径": " enum('在售', '停售', '已到期')"}
                    }
                },
                "客户购买理财产品记录表": {
                    "表名称": "customer_products",
                    "字段": {
                        "record_id": {"含义": "记录唯一标识", "计算口径": "自增整数"},
                        "customer_id": {"含义": "客户ID", "计算口径": "外键关联customers表customer_id"},
                        "product_id": {"含义": "产品ID", "计算口径": "外键关联financial_products表product_id"},
                        "investment_amount": {"含义": "投资金额", "计算口径": " decimal(15,2)"},
                        "purchase_date": {"含义": "购买日期", "计算口径": " datetime"},
                        "status": {"含义": "状态", "计算口径": " enum('持有', '赎回', '到期')"}
                    }
                }
            }
        }
    }
    return knowledge



# 5.工具注册管理
tools = [bank_knowledge_tool,execute_sql_tool]

# 6.绑定工具到模型
bind_qwen = qwen.bind(
    functions=[convert_to_openai_function(t) for t in tools]
)

# 7.创建代理
agent = (
        {
            "input": lambda x: x["input"],
            "agent_scratchpad": lambda x: format_to_openai_function_messages(
                x["intermediate_steps"]
            ),

        }
        | prompt
        | bind_qwen
        | OpenAIFunctionsAgentOutputParser()
)

# 8.用户输入和中间步骤初始化
user_questy = "查询存款总额前 10 的客户?"
intermediate_steps = []

# 9.代理执行循环 这正是ReAct范式（Reasoning-Acting循环）的具体实现，标准化Agent通过AgentType.ZERO_SHOT_REACT_DESCRIPTION参数封装该逻辑。

while True:
    output = agent.invoke(
        {
            "input": user_questy,
            "intermediate_steps": intermediate_steps
        }
    )

    if isinstance(output, AgentFinish):
        final_result = output.return_values["output"]
        break
    else:
        print(f"工具名称:{output.tool}")
        print(f"工具输入：{output.tool_input}")
        tool = {"bank_knowledge_tool": bank_knowledge_tool,"execute_sql_tool":execute_sql_tool}[output.tool]
        tool_run = tool.run(output.tool_input)
        intermediate_steps.append((output, tool_run))
print(final_result)